//
// Created by alax on 4/15/17.
//

#include "testFunctions.h"

FunctionPass *llvm::createX86Rerandomization() {
    return new findLRinstructs();
}
char findLRinstructs::ID = 0;

static inline bool replaceRBPwithRSP(MachineBasicBlock::iterator &BBI, int stackSize){
    int baseRegIndex=findMemOprandFirstIndex(*BBI);
    if(baseRegIndex>=0){
        MachineOperand &p=BBI->getOperand(baseRegIndex + X86::AddrBaseReg);
        if(p.isReg() && p.getReg()==X86::RBP){
            int64_t dis=BBI->getOperand(baseRegIndex + X86::AddrDisp).getImm();
            if(dis>0){
                errs()<<"RBP offset > 0 ::";
                BBI->print(errs());
            }
            dis+= stackSize - 8;
            BBI->getOperand(baseRegIndex + X86::AddrDisp).setImm(dis);
            BBI->getOperand(baseRegIndex + X86::AddrBaseReg).setReg(X86::RSP);
            return true;
        } else{// if base register is not RBP
//            BBI->print(errs());
            return false;
        }
    }
    return false;
}

bool findLRinstructs::processLEAinstr(MachineBasicBlock::iterator &I) {
    MachineOperand &baseR=I->getOperand(1);
    if(baseR.isReg() && baseR.getReg() == X86::RSP){
        baseR.setReg(X86::RBP);
        return true;
    }
    return false;
}
bool findLRinstructs::processBasicBlock(MachineFunction &MF,MachineFunction::iterator &MFI) {
    bool isModified=false;
    unsigned BBinstrNum = 0;
    for(MachineBasicBlock::iterator BBI = MFI->begin(),BBE = MFI->end(); BBI!=BBE;BBI++){
        BBinstrNum++;

        if(isLEA(BBI->getOpcode())){//
           //BBI->print(errs());
            processLEAinstr(BBI);
        }
    }
    return isModified;
}


bool findLRinstructs::runOnMachineFunction(MachineFunction &MF){
    //unsigned instructionNum=0;
    const MachineFrameInfo &MFI = MF.getFrameInfo();
    TII = MF.getSubtarget().getInstrInfo();
    TRI = MF.getSubtarget().getRegisterInfo();
    bool modified = false;
    curFuncStackSize = MFI.getStackSize();

    if(MF.getName()=="main") return false;
    if(MF.getName()=="fillPagetable") return false;
    //errs()<<"\n functionName:"<<MF.getName()<<" stack size "<<curFuncStackSize<<"\n";

    //handle each basic block
    for(MachineFunction::iterator I = MF.begin(), E=MF.end();I!=E;I++){//const_iterator is the iterator of BB
        modified |= processBasicBlock(MF,I);
    }
    return  modified;
}

static RegisterPass<findLRinstructs> X("findLRinstruction","find Load and store instruction");
